import struct
from typing import List, Optional, Tuple, Dict, Any
from .sub_array import SubArray
from utils.parser import Parser
from storage.data_item import DataItem

class Node:
    """B+树节点
    
    Node结构如下：
    [LeafFlag][KeyNumber][SiblingUid]
    [Son0][Key0][Son1][Key1]...[SonN][KeyN]
    """
    
    # 节点结构常量定义
    # 叶子节点标志
    IS_LEAF_OFFSET = 0
    # 无键标记
    NO_KEYS_OFFSET = IS_LEAF_OFFSET + 1
    # 兄弟标志
    SIBLING_OFFSET = NO_KEYS_OFFSET + 2
    # 节点头大小
    NODE_HEADER_SIZE = SIBLING_OFFSET + 8
    
    BALANCE_NUMBER = 32
    NODE_SIZE = NODE_HEADER_SIZE + (2*8)*(BALANCE_NUMBER*2+2)
    
    def __init__(self):
        """初始化节点"""
        self.tree = None  # BPlusTree实例
        self.data_item = None  # DataItem实例
        self.raw = None  # SubArray实例
        self.uid = 0  # 节点唯一标识符
    
    @staticmethod
    def set_raw_is_leaf(raw: SubArray, is_leaf: bool) -> None:
        """设置节点是否为叶子节点"""
        if is_leaf:
            raw.raw[raw.start + Node.IS_LEAF_OFFSET] = 1
        else:
            raw.raw[raw.start + Node.IS_LEAF_OFFSET] = 0
    
    @staticmethod
    def get_raw_if_leaf(raw: SubArray) -> bool:
        """获取节点是否为叶子节点"""
        return raw.raw[raw.start + Node.IS_LEAF_OFFSET] == 1
    
    @staticmethod
    def set_raw_no_keys(raw: SubArray, no_keys: int) -> None:
        """设置节点键数量"""
        # 使用小端序存储short值
        no_keys_bytes = struct.pack('<H', no_keys)
        raw.raw[raw.start + Node.NO_KEYS_OFFSET:raw.start + Node.NO_KEYS_OFFSET + 2] = no_keys_bytes
    
    @staticmethod
    def get_raw_no_keys(raw: SubArray) -> int:
        """获取节点键数量"""
        # 使用小端序解析short值
        return struct.unpack('<H', bytes(raw.raw[raw.start + Node.NO_KEYS_OFFSET:raw.start + Node.NO_KEYS_OFFSET + 2]))[0]
    
    @staticmethod
    def set_raw_sibling(raw: SubArray, sibling: int) -> None:
        """设置兄弟节点UID"""
        sibling_bytes = Parser.long_to_bytes(sibling)
        raw.raw[raw.start + Node.SIBLING_OFFSET:raw.start + Node.SIBLING_OFFSET + 8] = sibling_bytes
    
    @staticmethod
    def get_raw_sibling(raw: SubArray) -> int:
        """获取兄弟节点UID"""
        return Parser.parse_long(bytes(raw.raw[raw.start + Node.SIBLING_OFFSET:raw.start + Node.SIBLING_OFFSET + 8]))
    
    @staticmethod
    def set_raw_kth_son(raw: SubArray, uid: int, kth: int) -> None:
        """设置第k个子节点UID"""
        offset = raw.start + Node.NODE_HEADER_SIZE + kth * (8 * 2)
        uid_bytes = Parser.long_to_bytes(uid)
        raw.raw[offset:offset + 8] = uid_bytes
    
    @staticmethod
    def get_raw_kth_son(raw: SubArray, kth: int) -> int:
        """获取第k个子节点UID"""
        offset = raw.start + Node.NODE_HEADER_SIZE + kth * (8 * 2)
        return Parser.parse_long(bytes(raw.raw[offset:offset + 8]))
    
    @staticmethod
    def set_raw_kth_key(raw: SubArray, key: int, kth: int) -> None:
        """设置第k个键值"""
        offset = raw.start + Node.NODE_HEADER_SIZE + kth * (8 * 2) + 8
        key_bytes = Parser.long_to_bytes(key)
        raw.raw[offset:offset + 8] = key_bytes
    
    @staticmethod
    def get_raw_kth_key(raw: SubArray, kth: int) -> int:
        """获取第k个键值"""
        offset = raw.start + Node.NODE_HEADER_SIZE + kth * (8 * 2) + 8
        return Parser.parse_long(bytes(raw.raw[offset:offset + 8]))
    
    @staticmethod
    def copy_raw_from_kth(from_raw: SubArray, to_raw: SubArray, kth: int) -> None:
        """从源节点的第k个位置开始复制到目标节点"""
        offset = from_raw.start + Node.NODE_HEADER_SIZE + kth * (8 * 2)
        length = from_raw.end - offset
        to_raw.raw[to_raw.start + Node.NODE_HEADER_SIZE:to_raw.start + Node.NODE_HEADER_SIZE + length] = \
            from_raw.raw[offset:offset + length]
    
    @staticmethod
    def shift_raw_kth(raw: SubArray, kth: int) -> None:
        """将第k个位置及之后的数据向后移动一个位置"""
        begin = raw.start + Node.NODE_HEADER_SIZE + (kth + 1) * (8 * 2)
        end = raw.start + Node.NODE_SIZE - 1
        for i in range(end, begin - 1, -1):
            if i - (8 * 2) >= raw.start:
                raw.raw[i] = raw.raw[i - (8 * 2)]
    
    @staticmethod
    def new_root_raw(left: int, right: int, key: int) -> bytes:
        """创建新的根节点"""
        raw = SubArray(bytearray(Node.NODE_SIZE), 0, Node.NODE_SIZE)
        
        Node.set_raw_is_leaf(raw, False)
        Node.set_raw_no_keys(raw, 2)
        Node.set_raw_sibling(raw, 0)
        Node.set_raw_kth_son(raw, left, 0)
        Node.set_raw_kth_key(raw, key, 0)
        Node.set_raw_kth_son(raw, right, 1)
        Node.set_raw_kth_key(raw, 9223372036854775807, 1)  # Long.MAX_VALUE
        
        return bytes(raw.raw)
    
    @staticmethod
    def new_nil_root_raw() -> bytes:
        """创建空的根节点"""
        raw = SubArray(bytearray(Node.NODE_SIZE), 0, Node.NODE_SIZE)
        
        Node.set_raw_is_leaf(raw, True)
        Node.set_raw_no_keys(raw, 0)
        Node.set_raw_sibling(raw, 0)
        
        return bytes(raw.raw)
    
    @staticmethod
    def load_node(b_tree, uid: int) -> 'Node':
        """加载节点"""
        di = b_tree.dm.read(uid)
        assert di is not None
        n = Node()
        n.tree = b_tree
        n.data_item = di
        n.raw = di.data()
        n.uid = uid
        return n
    
    def release(self) -> None:
        """释放节点资源"""
        self.data_item.release()
    
    def is_leaf(self) -> bool:
        """判断是否为叶子节点"""
        self.data_item.r_lock()  # 读锁
        try:
            return Node.get_raw_if_leaf(self.raw)
        finally:
            self.data_item.r_unlock()
    # 搜索下一个结果集，uid单个
    class SearchNextRes:
        """搜索下一个节点的结果"""
        def __init__(self):
            self.uid = 0  # 子节点UID
            self.sibling_uid = 0  # 兄弟节点UID
    # 范围查询结果集
    class LeafSearchRangeRes:
        """叶子节点范围查询结果"""
        def __init__(self):
            self.uids = []  # 匹配的UID列表
            self.sibling_uid = 0  # 兄弟节点UID
    # 找找到下一个节点【等于key的】
    def search_next(self, key: int) -> 'Node.SearchNextRes':
        """查找下一个节点"""
        self.data_item.r_lock()  # 读锁
        try:
            # 初始化中间结果
            res = self.SearchNextRes()
            # 获取本节点key数量
            no_keys = Node.get_raw_no_keys(self.raw)
            
            # 遍历所有键，找到第一个大于等于key的键
            for i in range(no_keys):
                # 第k个键值
                ik = Node.get_raw_kth_key(self.raw, i)
                if key < ik:
                    # 对应的第K个孩子，K孩子里的都比K_key小
                    res.uid = Node.get_raw_kth_son(self.raw, i)
                    res.sibling_uid = 0
                    return res
            
            # 如果所有键都小于key，返回兄弟节点
            res.uid = 0
            res.sibling_uid = Node.get_raw_sibling(self.raw)
            return res
        finally:
            self.data_item.r_unlock()
    

    
    def leaf_search_range(self, left_key: int, right_key: int) -> 'Node.LeafSearchRangeRes':
        """叶子节点范围查询"""
        self.data_item.r_lock()  # 读锁
        try:
            no_keys = Node.get_raw_no_keys(self.raw)
            
            # 找到第一个大于等于left_key的位置
            kth = 0
            while kth < no_keys:
                ik = Node.get_raw_kth_key(self.raw, kth)
                if ik >= left_key:
                    break
                kth += 1
            
            # 收集范围内的所有UID
            uids = []
            while kth < no_keys:
                ik = Node.get_raw_kth_key(self.raw, kth)
                if ik <= right_key:
                    uids.append(Node.get_raw_kth_son(self.raw, kth))
                    kth += 1
                else:
                    break
            
            # 如果已经到达节点末尾，设置兄弟节点UID
            sibling_uid = 0
            if kth == no_keys:
                sibling_uid = Node.get_raw_sibling(self.raw)
            
            res = self.LeafSearchRangeRes()
            res.uids = uids
            res.sibling_uid = sibling_uid
            return res
        finally:
            self.data_item.r_unlock()
    

    
    def insert_and_split(self, uid: int, key: int) -> 'Node.InsertAndSplitRes':
        """插入并分裂节点"""
        success = False
        err = None
        res = self.InsertAndSplitRes()
        
        self.data_item.before()  # 事务开始
        try:
            success = self.insert(uid, key)
            if not success:
                res.sibling_uid = Node.get_raw_sibling(self.raw)
                return res
            
            if self.need_split():
                try:
                    r = self.split()
                    res.new_son = r.new_son
                    res.new_key = r.new_key
                    return res
                except Exception as e:
                    err = e
                    raise e
            else:
                return res
        finally:
            if err is None and success:
                # 使用超级事务ID提交
                self.data_item.after(0)  # 假设0为超级事务ID
            else:
                self.data_item.un_before()
    
    def insert(self, uid: int, key: int) -> bool:
        """插入键值对
        uid：
        """
        no_keys = Node.get_raw_no_keys(self.raw)
        
        # 找到插入位置
        kth = 0
        # 遍历找到第一个大于等于Key的值
        while kth < no_keys:
            # 获得第k个key
            ik = Node.get_raw_kth_key(self.raw, kth)
            if ik < key:
                kth += 1
            else:
                break
        
        # 如果已经到达节点末尾且有兄弟节点，返回失败
        if kth == no_keys and Node.get_raw_sibling(self.raw) != 0:
            return False
        
        # 根据节点类型执行不同的插入逻辑
        if Node.get_raw_if_leaf(self.raw):
            # 叶子节点：直接插入键值对
            Node.shift_raw_kth(self.raw, kth)
            Node.set_raw_kth_key(self.raw, key, kth)
            Node.set_raw_kth_son(self.raw, uid, kth)
            Node.set_raw_no_keys(self.raw, no_keys + 1)
        else:
            # 内部节点：插入键并移动现有键
            kk = Node.get_raw_kth_key(self.raw, kth)
            Node.set_raw_kth_key(self.raw, key, kth)
            Node.shift_raw_kth(self.raw, kth + 1)
            Node.set_raw_kth_key(self.raw, kk, kth + 1)
            Node.set_raw_kth_son(self.raw, uid, kth + 1)
            Node.set_raw_no_keys(self.raw, no_keys + 1)
        
        return True
    
    def need_split(self) -> bool:
        """判断是否需要分裂"""
        return Node.BALANCE_NUMBER * 2 == Node.get_raw_no_keys(self.raw)
    
    class SplitRes:
        """分裂结果"""
        def __init__(self):
            self.new_son = 0  # 新子节点UID
            self.new_key = 0  # 新键
    class InsertAndSplitRes:
        """插入并分裂结果"""
        def __init__(self):
            self.sibling_uid = 0  # 兄弟节点UID
            self.new_son = 0  # 新子节点UID
            self.new_key = 0  # 新键
    def split(self) -> 'Node.SplitRes':
        """分裂节点"""
        # 创建新节点
        node_raw = SubArray(bytearray(Node.NODE_SIZE), 0, Node.NODE_SIZE)
        Node.set_raw_is_leaf(node_raw, Node.get_raw_if_leaf(self.raw))
        Node.set_raw_no_keys(node_raw, Node.BALANCE_NUMBER)
        Node.set_raw_sibling(node_raw, Node.get_raw_sibling(self.raw))
        
        # 复制数据到新节点
        Node.copy_raw_from_kth(self.raw, node_raw, Node.BALANCE_NUMBER)
        
        # 插入新节点
        son = self.tree.dm.insert(0, bytes(node_raw.raw))  # 假设0为超级事务ID
        
        # 更新当前节点
        Node.set_raw_no_keys(self.raw, Node.BALANCE_NUMBER)
        Node.set_raw_sibling(self.raw, son)
        
        # 返回分裂结果
        res = self.SplitRes()
        res.new_son = son
        res.new_key = Node.get_raw_kth_key(node_raw, 0)
        return res
    
    def __str__(self) -> str:
        """生成节点信息字符串"""
        sb = []
        sb.append(f"Is leaf: {Node.get_raw_if_leaf(self.raw)}\n")
        key_number = Node.get_raw_no_keys(self.raw)
        sb.append(f"KeyNumber: {key_number}\n")
        sb.append(f"sibling: {Node.get_raw_sibling(self.raw)}\n")
        
        for i in range(key_number):
            sb.append(f"son: {Node.get_raw_kth_son(self.raw, i)}, key: {Node.get_raw_kth_key(self.raw, i)}\n")
        
        return ''.join(sb)